use std::collections::HashMap;
use std::collections::HashSet;

type Graph = HashMap<String, Vec<String>>;

fn parse_input() -> Graph {
    let input = std::fs::read_to_string("./day12/input.txt").unwrap();
    
    let mut graph = HashMap::new();
    for line in input.lines() {
        let index = line.find('-').unwrap();
        let node1 = line[..index].to_string();
        let node2 = line[(index+1)..].to_string();

        graph.entry(node1.clone()).or_insert(Vec::new()).push(node2.clone());
        graph.entry(node2).or_insert(Vec::new()).push(node1);       
    }

    graph
}

#[allow(unused)]
fn part1(graph: &Graph) {

    fn is_small_cave(caves: &str) -> bool {
        caves.chars().all(|c| c.is_lowercase())
    }

    fn backtrack(graph: &Graph, state: &mut HashSet<String>, caves: &String) -> usize {        
        if caves == "end" {
            return 1;
        }

        let mut count = 0;
        for next_cave in graph.get(caves).unwrap() {
            if is_small_cave(&next_cave) {
                if state.contains(next_cave) {
                    continue;
                }
                state.insert(next_cave.clone());
                count += backtrack(graph, state, next_cave);
                state.remove(next_cave);
            } else {
                count += backtrack(graph, state, next_cave);
            }
        }

        count
    }

    let mut state = HashSet::new();
    state.insert("start".to_string());
    println!("Part1: {}", backtrack(graph, &mut state, &"start".to_string()));

}

#[allow(unused)]
fn part2(graph: &Graph) {

    fn is_small_cave(caves: &str) -> bool {
        caves.chars().all(|c| c.is_lowercase())
    }

    fn backtrack(graph: &Graph, state: &mut HashSet<String>, caves: &String, twice_cave: Option<(&String, bool)>) -> usize {        
        if caves == "end" {
            if let Some((_, is_twice)) = twice_cave {
                // A twice try must have twice cave
                if !is_twice {
                    return 0;
                }
            }
            return 1;
        }

        let mut count = 0;
        for next_cave in graph.get(caves).unwrap() {
            if is_small_cave(&next_cave) {
                if state.contains(next_cave) {
                    continue;
                }

                match twice_cave {
                    Some((cave, is_twice)) => {
                        state.insert(next_cave.clone());
                        
                        let is_twice = if cave == next_cave { true } else { is_twice };
                        count += backtrack(graph, state, next_cave, Some((cave, is_twice)));

                        state.remove(next_cave);
                    } 
                    None => {
                        state.insert(next_cave.clone());
                        // Try with no twice
                        count += backtrack(graph, state, next_cave, None);
                        state.remove(next_cave);
                        
                        // Try with twice
                        count += backtrack(graph, state, next_cave, Some((next_cave, false)));
                    }
                }
            } else {
                count += backtrack(graph, state, next_cave, twice_cave);
            }
        }

        count
    }

    let mut state = HashSet::new();
    state.insert("start".to_string());
    println!("Part2: {}", backtrack(graph, &mut state, &"start".to_string(), None));
}

fn main() {
    let graph = parse_input();
    part1(&graph);
    part2(&graph);
}
